#!/usr/bin/python

"""
Author: Anton Feenstra (feenstra@few.vu.nl)

Revision History:
File created on Nov 4, 2011

Template for dynamic programming exercises.

Copyright (c) 2011 Anton Feenstra
"""

'''
NOTE ON READABILITY:
 
 - script is organized into sections where one section is marke by:
   #-------------------------------------
   # Section description
   #-------------------------------------
 
 - in each section the upper level functions (function called early on) are
   at the top of the section and the low level functions are at the bottom 
   of the section
'''

import sys

# for command line options:
from optparse import OptionParser

# Alignment type enum
ALGN_GLOBAL = 0
ALGN_LOCAL = 1
ALGN_SEMI = 2

#-------------------------------------------------------------------------------
# Built-in exchange matrices.
#-------------------------------------------------------------------------------

identity = [
    [ 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,],
    [-1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,],
    [-1,-1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,],
    [-1,-1,-1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,],
    [-1,-1,-1,-1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,],
    [-1,-1,-1,-1,-1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,],
    [-1,-1,-1,-1,-1,-1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,],
    [-1,-1,-1,-1,-1,-1,-1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,],
    [-1,-1,-1,-1,-1,-1,-1,-1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,],
    [-1,-1,-1,-1,-1,-1,-1,-1,-1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,],
    [-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,],
    [-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,],
    [-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,],
    [-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,],
    [-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,],
    [-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,],
    [-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,],
    [-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 1,-1,-1,-1,-1,-1,-1,-1,-1,],
    [-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 1,-1,-1,-1,-1,-1,-1,-1,],
    [-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 1,-1,-1,-1,-1,-1,-1,],
    [-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 1,-1,-1,-1,-1,-1,],
    [-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 1,-1,-1,-1,-1,],
    [-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 1,-1,-1,-1,],
    [-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 1,-1,-1,],
    [-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 1,-1,],
    [-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 1,],
]

pam250 = [
    [ 2, 0,-2, 0, 0,-4, 1,-1,-1, 0,-1,-2,-1, 0, 0, 1, 0,-2, 1, 1, 0, 0,-6, 0,-3, 0,],
    [ 0, 2,-4, 3, 2,-5, 0, 1,-2, 0, 1,-3,-2, 2, 0,-1, 1,-1, 0, 0, 0,-2,-5, 0,-3, 2,],
    [-2,-4,12,-5,-5,-4,-3,-3,-2, 0,-5,-6,-5,-4, 0,-3,-5,-4, 0,-2, 0,-2,-8, 0, 0,-5,],
    [ 0, 3,-5, 4, 3,-6, 1, 1,-2, 0, 0,-4,-3, 2, 0,-1, 2,-1, 0, 0, 0,-2,-7, 0,-4, 3,],
    [ 0, 2,-5, 3, 4,-5, 0, 1,-2, 0, 0,-3,-2, 1, 0,-1, 2,-1, 0, 0, 0,-2,-7, 0,-4, 3,],
    [-4,-5,-4,-6,-5, 9,-5,-2, 1, 0,-5, 2, 0,-4, 0,-5,-5,-4,-3,-3, 0,-1, 0, 0, 7,-5,],
    [ 1, 0,-3, 1, 0,-5, 5,-2,-3, 0,-2,-4,-3, 0, 0,-1,-1,-3, 1, 0, 0,-1,-7, 0,-5,-1,],
    [-1, 1,-3, 1, 1,-2,-2, 6,-2, 0, 0,-2,-2, 2, 0, 0, 3, 2,-1,-1, 0,-2,-3, 0, 0, 2,],
    [-1,-2,-2,-2,-2, 1,-3,-2, 5, 0,-2, 2, 2,-2, 0,-2,-2,-2,-1, 0, 0, 4,-5, 0,-1,-2,],
    [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,],
    [-1, 1,-5, 0, 0,-5,-2, 0,-2, 0, 5,-3, 0, 1, 0,-1, 1, 3, 0, 0, 0,-2,-3, 0,-4, 0,],
    [-2,-3,-6,-4,-3, 2,-4,-2, 2, 0,-3, 6, 4,-3, 0,-3,-2,-3,-3,-2, 0, 2,-2, 0,-1,-3,],
    [-1,-2,-5,-3,-2, 0,-3,-2, 2, 0, 0, 4, 6,-2, 0,-2,-1, 0,-2,-1, 0, 2,-4, 0,-2,-2,],
    [ 0, 2,-4, 2, 1,-4, 0, 2,-2, 0, 1,-3,-2, 2, 0,-1, 1, 0, 1, 0, 0,-2,-4, 0,-2, 1,],
    [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,],
    [ 1,-1,-3,-1,-1,-5,-1, 0,-2, 0,-1,-3,-2,-1, 0, 6, 0, 0, 1, 0, 0,-1,-6, 0,-5, 0,],
    [ 0, 1,-5, 2, 2,-5,-1, 3,-2, 0, 1,-2,-1, 1, 0, 0, 4, 1,-1,-1, 0,-2,-5, 0,-4, 3,],
    [-2,-1,-4,-1,-1,-4,-3, 2,-2, 0, 3,-3, 0, 0, 0, 0, 1, 6, 0,-1, 0,-2, 2, 0,-4, 0,],
    [ 1, 0, 0, 0, 0,-3, 1,-1,-1, 0, 0,-3,-2, 1, 0, 1,-1, 0, 2, 1, 0,-1,-2, 0,-3, 0,],
    [ 1, 0,-2, 0, 0,-3, 0,-1, 0, 0, 0,-2,-1, 0, 0, 0,-1,-1, 1, 3, 0, 0,-5, 0,-3,-1,],
    [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,],
    [ 0,-2,-2,-2,-2,-1,-1,-2, 4, 0,-2, 2, 2,-2, 0,-1,-2,-2,-1, 0, 0, 4,-6, 0,-2,-2,],
    [-6,-5,-8,-7,-7, 0,-7,-3,-5, 0,-3,-2,-4,-4, 0,-6,-5, 2,-2,-5, 0,-6,17, 0, 0,-6,],
    [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,],
    [-3,-3, 0,-4,-4, 7,-5, 0,-1, 0,-4,-1,-2,-2, 0,-5,-4,-4,-3,-3, 0,-2, 0, 0,10,-4,],
    [ 0, 2,-5, 3, 3,-5,-1, 2,-2, 0, 0,-3,-2, 1, 0, 0, 3, 0, 0,-1, 0,-2,-6, 0,-4, 3,],
]

blosum62 = [
    [ 4,-2, 0,-2,-1,-2, 0,-2,-1, 0,-1,-1,-1,-2, 0,-1,-1,-1, 1, 0, 0, 0,-3, 0,-2,-1,],
    [-2, 4,-3, 4, 1,-3,-1, 0,-3, 0, 0,-4,-3, 3, 0,-2, 0,-1, 0,-1, 0,-3,-4,-1,-3, 1,],
    [ 0,-3, 9,-3,-4,-2,-3,-3,-1, 0,-3,-1,-1,-3, 0,-3,-3,-3,-1,-1, 0,-1,-2,-2,-2,-3,],
    [-2, 4,-3, 6, 2,-3,-1,-1,-3, 0,-1,-4,-3, 1, 0,-1, 0,-2, 0,-1, 0,-3,-4,-1,-3, 1,],
    [-1, 1,-4, 2, 5,-3,-2, 0,-3, 0, 1,-3,-2, 0, 0,-1, 2, 0, 0,-1, 0,-2,-3,-1,-2, 4,],
    [-2,-3,-2,-3,-3, 6,-3,-1, 0, 0,-3, 0, 0,-3, 0,-4,-3,-3,-2,-2, 0,-1, 1,-1, 3,-3,],
    [ 0,-1,-3,-1,-2,-3, 6,-2,-4, 0,-2,-4,-3, 0, 0,-2,-2,-2, 0,-2, 0,-3,-2,-1,-3,-2,],
    [-2, 0,-3,-1, 0,-1,-2, 8,-3, 0,-1,-3,-2, 1, 0,-2, 0, 0,-1,-2, 0,-3,-2,-1, 2, 0,],
    [-1,-3,-1,-3,-3, 0,-4,-3, 4, 0,-3, 2, 1,-3, 0,-3,-3,-3,-2,-1, 0, 3,-3,-1,-1,-3,],
    [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,],
    [-1, 0,-3,-1, 1,-3,-2,-1,-3, 0, 5,-2,-1, 0, 0,-1, 1, 2, 0,-1, 0,-2,-3,-1,-2, 1,],
    [-1,-4,-1,-4,-3, 0,-4,-3, 2, 0,-2, 4, 2,-3, 0,-3,-2,-2,-2,-1, 0, 1,-2,-1,-1,-3,],
    [-1,-3,-1,-3,-2, 0,-3,-2, 1, 0,-1, 2, 5,-2, 0,-2, 0,-1,-1,-1, 0, 1,-1,-1,-1,-1,],
    [-2, 3,-3, 1, 0,-3, 0, 1,-3, 0, 0,-3,-2, 6, 0,-2, 0, 0, 1, 0, 0,-3,-4,-1,-2, 0,],
    [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,],
    [-1,-2,-3,-1,-1,-4,-2,-2,-3, 0,-1,-3,-2,-2, 0, 7,-1,-2,-1,-1, 0,-2,-4,-2,-3,-1,],
    [-1, 0,-3, 0, 2,-3,-2, 0,-3, 0, 1,-2, 0, 0, 0,-1, 5, 1, 0,-1, 0,-2,-2,-1,-1, 3,],
    [-1,-1,-3,-2, 0,-3,-2, 0,-3, 0, 2,-2,-1, 0, 0,-2, 1, 5,-1,-1, 0,-3,-3,-1,-2, 0,],
    [ 1, 0,-1, 0, 0,-2, 0,-1,-2, 0, 0,-2,-1, 1, 0,-1, 0,-1, 4, 1, 0,-2,-3, 0,-2, 0,],
    [ 0,-1,-1,-1,-1,-2,-2,-2,-1, 0,-1,-1,-1, 0, 0,-1,-1,-1, 1, 5, 0, 0,-2, 0,-2,-1,],
    [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,],
    [ 0,-3,-1,-3,-2,-1,-3,-3, 3, 0,-2, 1, 1,-3, 0,-2,-2,-3,-2, 0, 0, 4,-3,-1,-1,-2,],
    [-3,-4,-2,-4,-3, 1,-2,-2,-3, 0,-3,-2,-1,-4, 0,-4,-2,-3,-3,-2, 0,-3,11,-2, 2,-3,],
    [ 0,-1,-2,-1,-1,-1,-1,-1,-1, 0,-1,-1,-1,-1, 0,-2,-1,-1, 0, 0, 0,-1,-2,-1,-1,-1,],
    [-2,-3,-2,-3,-2, 3,-3, 2,-1, 0,-2,-1,-1,-2, 0,-3,-1,-2,-2,-2, 0,-1, 2,-1, 7,-2,],
    [-1, 1,-3, 1, 4,-3,-2, 0,-3, 0, 1,-3,-1, 0, 0,-1, 3, 0, 0,-1, 0,-2,-3,-1,-2, 4,],
]

#-------------------------------------------------------------------------------
# Main function, command line args, sequence helpers, print helpers
# (i.e. the given template + little extra)
#-------------------------------------------------------------------------------

def main():
    # get command line options
    options = parse_commandline()

    # set substitution matrix:
    if options.exchange_matrix == "pam250":
        exchangeMatrix = pam250
    elif options.exchange_matrix == "blosum62":
        exchangeMatrix = blosum62
    elif options.exchange_matrix == "identity":
        exchangeMatrix = identity
    else:
        print "unknown exchange matrix", options.exchange_matrix
        exit(-1)
    
    # read sequences from fasta file, and catch error reading file
    try:
        sequences = readSequences(open(options.fasta))
    except IOError:
        print "ERROR: cannot open or read fasta input file:", options.fasta
        exit(-1)

    for seq in sequences:
        print seq
    
    # scoring matrix, pointer matrix and alignment
    sm = []
    pm = []
    algn = []
    
    # call alignment routine(s):
    if options.align_global:
        (algn, sm, pm) = align(
            sequences, exchangeMatrix, options.gap_penalty, ALGN_GLOBAL)
    elif options.align_local:
        (algn, sm, pm) = align(
            sequences, exchangeMatrix, options.gap_penalty, ALGN_LOCAL)
    elif options.align_semiglobal:
        (algn, sm, pm) = align(
            sequences, exchangeMatrix, options.gap_penalty, ALGN_SEMI)
    else:
        print "BUG! this should not happen."
        exit(-1)
        
    # Print scoring matrix
    if options.verbose:        
        # First append the the amino acids of the sequences to the matrix
        y_seq = "-" + sequences[0].Sequence
        x_seq = " -" + sequences[1].Sequence        
        for i in range(len(sm)):
            sm[i].insert(0,y_seq[i])
        sm.insert(0, x_seq)
        
        print_matrix(sm)
        
    print_algn(algn)


def parse_commandline():
    usage = "%prog <fasta> [options]"
    version = "0.0a"
    description = \
        "%prog aligns two sequences."
    epilog = \
        "Copyright (c) 2011 K. Anton Feenstra -- "\
        "feenstra@few.vu.nl -- www.few.vu.nl/~feenstra"
    parser = OptionParser(usage=usage, description=description,
                          version="%prog "+version, epilog=epilog)

    # sequence/alignment options:
    parser.add_option("-f", "--fasta",  dest="fasta", metavar="<file>",
                     help="input alignment file (fasta)")
    parser.set_defaults(fasta=None)
    
    parser.add_option("-e", "",  dest="exchange_matrix", 
                     help="Exchange matrix: pam250, blosum62 or identity (default is %default)")
    parser.set_defaults(exchange_matrix="pam250")
    
    parser.add_option("-v", "",  dest="verbose", action="store_true", 
                     help="Prints out the scoring matrix")    
    parser.set_defaults(verbose=False)
    
    parser.add_option("-l", "",  dest="align_local",  action="store_true",
                     help="align local")
    parser.set_defaults(align_local=False)
    
    parser.add_option("-g", "",  dest="align_global", action="store_true",
                     help="align global")
    parser.set_defaults(align_global=False)
    
    parser.add_option("-s", "",  dest="align_semiglobal", action="store_true",
                     help="align semi-global")
    parser.set_defaults(align_semiglobal=False)
    
    parser.add_option("-p", "",  dest="gap_penalty", type="int",
                     help="Gap penalty (default is %default)")
    parser.set_defaults(gap_penalty=2)
    
    # get the options:
    (options, args) = parser.parse_args()

    if not options.fasta:
        # check if we have an option left (to be used as input filename):
        if args:
            options.fasta = args.pop()
        else:
            print "Need at least an input file (fasta)"
            print ""
            parser.print_help()
            print ""
            print "ERROR: no input file given"
            exit(-1)

    # check alignment type:
    align_options = [options.align_local, options.align_global, options.align_semiglobal]
    # check if at least one alignment option was true, else choose global
    if align_options.count(True)==0:
        print "No alignment type given, using Global"
        options.align_global=True
    # check if not more than one alignment option was true, else error and exit 
    if align_options.count(True)>1:
        print "ERROR: multiple alignment types chosen"
        exit(-1)

    # check for any leftover command line arguments:
    if len(args):
        print "ignoring additional arguments "+str(args)
    
    # clean up (recommended):
    del(parser)
    return options


class Sequence:
    """Stores a sequence object"""
    
    def __init__(self, Label="", Sequence="" ):
        """Initialize a new Sequence object

        Label -- identifier of sequence (text)
        Sequence -- sequence string in single-letter alphabet
        """
        self.Label       = Label
        self.Sequence    = Sequence

    # this makes that you can do 'print sequence' and get nice output:
    def __str__(self):
        """Return string representation of a Sequence object"""
        # newline-delimited values of all the attributes
        return ">%s\n%s" % (self.Label, self.Sequence)


def readSequences(lines):
    """Return Sequences object
    
    lines -- list of lines or any object that behaves like it
    
    This routine parses a fasta file and returns a list of Sequence objects
    containing the sequences with label and sequence data set
    """
    seqs = []
    label = None
    seq_lines = []
    for line in lines:
        line = line.strip()      # strip off white space
        if not line:             # skip empty lines
            continue
        if line.startswith(';'): # ignore comment lines
            continue
        # check for start of next sequence:
        if line.startswith('>'): # label line
            # first, store the previous sequence if we had one:
            if seq_lines:
                seqs.append(Sequence(label, ''.join(seq_lines)))
                seq_lines = []
            # get the label (name) for the next sequence
            label = line[1:].strip()
        else:
            # collect all lines with sequence information for this sequence:
            seq_lines.append(line)
    # take care of the last sequence in the file
    seqs.append(Sequence(label, ''.join(seq_lines)))
    return seqs
    
    
'''
Prints 2D matrix with each element justified to 5 places to the right
'''    
def print_matrix(m):
    print
    for i in range(len(m)):
        for j in range(len(m[0])):
            print str(m[i][j]).rjust(5) ,
        print
        

'''
Prints the final alignment in the format:
A B - C
  |   |
- B G C H
'''        
def print_algn(algn):
    y_algn = algn[0]
    x_algn = algn[1]
    match = algn[2]
    score = algn[3]
        
    match_str = ''
    for i in range(len(match)):
        if match[i]:
            match_str += '|'
        else:
            match_str += ' '
            
    print
    print ' '.join(y_algn)
    print ' '.join(match_str)
    print ' '.join(x_algn)
    print 'score = ' + str(score)

#-------------------------------------------------------------------------------
# Alignment
# (the work which needed to be done - all new)
#-------------------------------------------------------------------------------

'''
Do pairwise alignment using DP.
What type of alignment is done (local, global or semi-global) is decided based
on the algn_type parameter.
'''
def align(sequences, exchange_matrix, penalty, algn_type):
    
    # Based on the alignment type choose:
    #    scoring matrix initialization function
    #    pointer matrix initialization function
    #    default maximum value for scoring matrix calculation
    #    traceback function
    
    init_sm = None
    init_pm = None
    default_max = None
    traceback = None

    if algn_type == ALGN_GLOBAL:        
        init_sm = init_sm_global
        init_pm = init_pm_global
        # Setting the default maximum really low enables us to override
        # the maximum with one of the values calculated from the diagonal, 
        # left or upper cells of diagonal matrix.
        default_max = -sys.maxint
        traceback = traceback_global
                
    elif algn_type == ALGN_LOCAL:    
        init_sm = init_sm_local
        init_pm = init_pm_local
        # In the local alignment we do not want values lower than 0
        # in the scoring matrix
        default_max = 0
        traceback = traceback_local
            
    elif algn_type == ALGN_SEMI:                
        init_sm = init_sm_semi
        init_pm = init_pm_semi
        default_max = -sys.maxint
        traceback = traceback_semi

    # Once we have the functions for the chosen alignment type
    # do the alignment

    (sm, pm) = create_score_matrix(
        sequences, exchange_matrix, penalty, init_sm, init_pm, default_max)    
    algn = traceback(sm, pm, sequences)
    
    return algn, sm, pm


'''
Returns: (scoring matrix, pointer matrix)
init_func: function which initializes the scoring matrix 
    (specific to the chosen alignment type - global, local or semi-global)
formula: formula for calculating each cell in the scoring matrix
    (specific to the chosen alignment type - global, local or semi-global)
'''
def create_score_matrix(
    sequences, exchange_matrix, penalty, init_sm, init_pm, default_max):
    
    y_seq = sequences[0].Sequence
    x_seq = sequences[1].Sequence
    
    # + 1 is for the initial gap penalty
    numOfRows = len(y_seq) + 1
    numOfCols = len(x_seq) + 1
    
    sm = create_matrix(numOfRows, numOfCols, 0)
    init_sm(sm, penalty)
    pm = create_matrix(numOfRows, numOfCols, (0,0))
    init_pm(pm)
    
    # Score matrix calculation
    for i in range(1, numOfRows):
        for j in range(1, numOfCols):
            
            diagonal = sm[i-1][j-1]
            left = sm[i][j-1]
            up = sm[i-1][j]
            
            y_aa = y_seq[i-1]
            x_aa = x_seq[j-1]
            exchange_value = exchange_matrix[ord(y_aa)-ord('A')][ord(x_aa)-ord('A')]
            
            (sm[i][j], pm[i][j]) = sm_formula(
                diagonal, left, up, penalty, exchange_value, default_max)
    
    return sm, pm
        

'''
Introduce the initial gap penalty into the scoring matrix.
'''
def init_sm_global(sm, penalty):
    rows = len(sm)
    cols = len(sm[0])
    
    for i in range(rows):
        for j in range(cols):
            if i == 0:
                sm[i][j] = -(penalty * j)
            if j == 0:
                sm[i][j] = -(penalty * i)        


'''
Do nothing.
In local alignment there is no need to introduce the initial gap penalty.
'''
def init_sm_local(sm, penalty):
    pass
                

'''
Do nothing.
In semi-global alignment there is no need to introduce the initial gap penalty.
'''
def init_sm_semi(sm, penalty):
    pass    


def init_pm_global(pm):
    init_pm_common(pm)


def init_pm_semi(pm):
    init_pm_common(pm)


'''
Initializes the first row and first column of the pointer matrix.
We can reach position [0,j] only from [0,j-1] (similarly with the first row).
This initialization is necessary for global and semi-global alignements where
we want to see all the residues in the final alignment.
'''
def init_pm_common(pm):
    numOfRows = len(pm)
    numOfCols = len(pm[0])
    for i in range(1,numOfRows):
        pm[i][0] = (-1,0)
    for j in range(1,numOfCols):
        pm[0][j] = (0,-1)
    

'''
Do nothing.
In local alignment there is no need to initialize the pointer matrix since
we do not want the alignment to continue once we hit 
the first column or first row.
'''
def init_pm_local(pm):
    pass
                                

'''
Max formula for calculating the value of the scoring matrix.
Returns: (maximum, direction from which we calculated the maximum)
'''
def sm_formula(diagonal, left, up, penalty, exchange_value, default_max):
    pointer = (0,0)
    mymax = default_max

    if diagonal + exchange_value > mymax:
        mymax = diagonal + exchange_value
        pointer = (-1,-1)
    
    if left - penalty > mymax:
        mymax = left - penalty
        pointer = (0,-1)
        
    if up - penalty > mymax:
        mymax = up - penalty
        pointer = (-1,0)
        
    return (mymax, pointer)


'''
Creates 2D matrix initialised with the given item. 
'''
def create_matrix(rows, cols, item):
    retval = []
    for i in range(rows):
        retval.append([])
        for j in range(cols):
            retval[i].append(item)
    return retval


def traceback_global(sm, pm, sequences):
    (start_i, start_j) = tb_start_global(sm)
    score = sm[start_i][start_j]
    y_algn, x_algn, match = traceback_common(pm, (start_i, start_j), sequences) 
    return y_algn, x_algn, match, score


def traceback_local(sm, pm, sequences):
    (start_i, start_j) = tb_start_local(sm)
    score = sm[start_i][start_j]
    y_algn, x_algn, match = traceback_common(pm, (start_i, start_j), sequences) 
    return y_algn, x_algn, match, score


'''
In the semi-global alignment the common traceback method needs to be extended.
We need to append the trailing gaps which were not penalized.
'''
def traceback_semi(sm, pm, sequences):
    (start_i, start_j) = tb_start_semi(sm)
    (y_algn, x_algn, match) = traceback_common(pm, (start_i, start_j), sequences)
    
    lastRow = len(sm) - 1
    lastCol = len(sm[0]) - 1
    y_seq = sequences[0].Sequence
    x_seq = sequences[1].Sequence
    
    # When we started the traceback at the last row add the rest of the row 
    # to the alignment
    if (start_i == lastRow) and (start_j != lastCol):
        for j in range(start_j + 1, lastCol + 1):
            y_algn.append('-')
            x_algn.append(x_seq[j-1])
            match.append(False)
    
    # When we started the traceback at the last column add the rest 
    # of the column to the alignment        
    if (start_j == lastCol) and (start_i != lastRow):
        for i in range(start_i + 1, lastRow + 1):
            y_algn.append(y_seq[i-1])
            x_algn.append('-')
            match.append(False)
            
    score = sm[start_i][start_j]
    return y_algn, x_algn, match, score 


'''
Based on the given start traces back through the pointer matrix 
and constructs the final alignment.
Returns: ([y sequence with gaps], [x sequence with gaps], [True if x and y matches])
'''
def traceback_common(pm, tb_start, sequences):
    y_algn = []
    x_algn = []
    match = []
    
    y_seq = sequences[0].Sequence
    x_seq = sequences[1].Sequence
    
    (current_i, current_j) = tb_start
    (direction_i, direction_j) = pm[current_i][current_j]
    
    # We stop the tracebeck when we reach a position from which we can not 
    # navigate to another cell.
    # For global and semi-global alignment only the upper left corner of the
    # pointer matrix can have no further direction (i.e. 0,0).
    # Howevere, the pointer matrix of local alignment will have no further 
    # direction anywhere where the scoring matrix is 0.
    while (direction_i, direction_j) != (0,0):        
        
        y_aa = y_seq[current_i - 1]
        x_aa = x_seq[current_j - 1]
        
        # We came from diagonal cell
        if (direction_i, direction_j) == (-1,-1):
            y_algn.append(y_aa)
            x_algn.append(x_aa)
            if y_aa == x_aa:
                match.append(True)
            else:
                match.append(False)
                
        # We came from left cell
        elif (direction_i, direction_j) == (0,-1):
            y_algn.append('-')
            x_algn.append(x_aa)
            match.append(False)
            
        # We came from upper cell
        elif (direction_i, direction_j) == (-1,0):
            y_algn.append(y_aa)
            x_algn.append('-')
            match.append(False)
            
        # Invalid direction
        else:
            print 'Invalid direction!'
            exit(-1)
            
        (current_i, current_j) = (current_i + direction_i, current_j + direction_j)
        (direction_i, direction_j) = pm[current_i][current_j]
    
    # Since we traced back from the bottom right corner of the matrix we need
    # to reverse the final alignment
    y_algn.reverse()
    x_algn.reverse()
    match.reverse()
    return y_algn, x_algn, match 
             

'''
When doing the global alignment we start the traceback_common in bottom right corner
of the scoring matrix
'''
def tb_start_global(sm):
    i = len(sm) - 1
    j = len(sm[0]) - 1
    return i,j


'''
When doing the local alignment we start the traceback_common at the maximum value 
found anywhere in the matrix
'''
def tb_start_local(sm):
    mymax = -sys.maxint
    numOfRows = len(sm)
    numOfCols = len(sm[0])
    start_i = 0
    start_j = 0
    matrix_size_diff = abs(numOfRows - numOfCols)
    
    for i in range(numOfRows):
        for j in range(numOfCols):
            (mymax, start_i, start_j) = test_max(
                sm[i][j], i, j, mymax, start_i, start_j, matrix_size_diff)        
        
    assert (start_i, start_j) != (0,0)
    return start_i, start_j


'''
When doing the semi-global alignment we start the traceback_common from the maximum
value found in the last column or last row
''' 
def tb_start_semi(sm):
    mymax = -sys.maxint
    numOfRows = len(sm)
    numOfCols = len(sm[0])
    start_i = 0
    start_j = 0
    matrix_size_diff = abs(numOfRows - numOfCols)
    
    # Search the last column for the max
    for i in range(numOfRows):
        (mymax, start_i, start_j) = test_max(
            sm[i][numOfCols - 1], i, numOfCols - 1, mymax, start_i, start_j,
            matrix_size_diff)
    
    # Search the last row for the max
    for j in range(numOfCols):
        (mymax, start_i, start_j) = test_max(
            sm[numOfRows - 1][j], numOfRows - 1, j, mymax, start_i, start_j,
            matrix_size_diff)
        
    assert (start_i, start_j) != (0,0)
    return start_i, start_j


'''
 Out of two given values returns the larger one together with its 
 i and j coordinates.
 If the values are equal returns the max which is closer to the diagonal 
 starting at the bottom right corner.
 Returns: (max, i, j)
'''
def test_max(candidate, i, j, current_max, current_i, current_j, matrix_size_diff):
    if candidate > current_max:
        return candidate, i, j
    
    if (candidate == current_max) \
        and (abs(abs(i - j) - matrix_size_diff) 
             < abs(abs(current_i - current_j) - matrix_size_diff)):
        return candidate, i, j
    
    return current_max, current_i, current_j
    
#-------------------------------------------------------------------------------
# Body of the script
#-------------------------------------------------------------------------------

if __name__ == "__main__":
    main()
